import os

import cv2
import torch
from torchvision import transforms

from mnist.ResidualCNNTrain_withCUDA import ResidualNeuralNetwork, device, transform

if __name__ == '__main__':
    # 加载模型
    model = ResidualNeuralNetwork()
    model.load_state_dict(torch.load("models/best.pth"))
    model.to(device)
    # 归一化,使得每次预测的结果一致
    model.eval()

    # 用于归一化的参数
    normalize = transforms.Normalize((0.1307,), (0.3081,))
    # 先转换到tensor数据
    transToTensor = transforms.ToTensor()

    # 文件夹
    folderPath = "testImg"
    # 遍历文件夹并读取图片
    allFiles = os.listdir(folderPath)
    for file in allFiles:
        # 分离文件名和后缀
        name, suffix = os.path.splitext(file)
        if suffix.lower() == ".jpg" or suffix.lower() == ".png":
            filePath = os.path.join(folderPath, file)
            # 读取图片
            img = cv2.imread(filePath, cv2.IMREAD_GRAYSCALE)
            # 压缩图片到训练size
            img = cv2.resize(img, (28, 28))
            # 反转颜色
            img = 255 - img
            # 显示图片
            cv2.imshow("test", img)
            cv2.waitKey(0)

            # 将图片转换到tensor
            # imgTensor = transToTensor(img)
            # 归一化
            testData = transform(img)
            # 用来观察数据的测试数据
            tempImageData = testData.numpy()
            # 数据降维
            # testData = testData.view(-1, 28 * 28)
            # 数据索引0处增加1个维度1
            testData = testData.unsqueeze(dim=0)
            # 测试
            testData = testData.to(device)
            r = model(testData)
            val = torch.max(r, dim=1)
            print(file, val.indices.item())
